1929F - Sasha and the Wedding Binary Search Tree - CodeForces Solution

#include <bits/stdc++.h>

typedef long long valueType;
typedef std::vector<valueType> ValueVector;
typedef std::vector<ValueVector> ValueMatrix;
typedef std::vector<ValueMatrix> ValueCube;
typedef std::string string;
typedef std::vector<string> StringVector;
typedef std::vector<bool> bitset;
typedef std::vector<bitset> BitMatrix;
typedef std::pair<valueType, valueType> ValuePair;
typedef std::vector<ValuePair> PairVector;
typedef std::vector<PairVector> PairMatrix;
typedef std::tuple<valueType, valueType, valueType> ValueTuple;
typedef std::vector<ValueTuple> TupleVector;
typedef std::vector<TupleVector> TupleMatrix;
typedef std::set<valueType> ValueSet;

namespace MODINT_WITH_FIXED_MOD {
    constexpr valueType MOD = 998244353;

    template<typename T1, typename T2>
    void Inc(T1 &a, T2 b) {
        a = a + b;

        if (a >= MOD)
            a -= MOD;
    }

    template<typename T1, typename T2>
    void Dec(T1 &a, T2 b) {
        a = a - b;

        if (a < 0)
            a += MOD;
    }

    template<typename T1, typename T2>
    T1 sum(T1 a, T2 b) {
        return a + b >= MOD ? a + b - MOD : a + b;
    }

    template<typename T1, typename T2>
    T1 sub(T1 a, T2 b) {
        return a - b < 0 ? a - b + MOD : a - b;
    }

    template<typename T1, typename T2>
    T1 mul(T1 a, T2 b) {
        return (long long) a * b % MOD;
    }

    template<typename T1, typename T2>
    void Mul(T1 &a, T2 b) {
        a = (long long) a * b % MOD;
    }

    template<typename T1, typename T2>
    T1 pow(T1 a, T2 b) {
        T1 result = 1;

        while (b > 0) {
            if (b & 1)
                Mul(result, a);

            Mul(a, a);
            b = b >> 1;
        }

        return result;
    }
} // namespace MODINT_WITH_FIXED_MOD

using namespace MODINT_WITH_FIXED_MOD;

class BinomialCoefficient {
private:
    valueType N;
    ValueVector Fact, InvFact;

public:
    BinomialCoefficient() = default;

    explicit BinomialCoefficient(valueType n) : N(n), Fact(n + 1, 1), InvFact(n + 1, 1) {
        for (valueType i = 1; i <= N; ++i)
            Fact[i] = mul(Fact[i - 1], i);

        InvFact[N] = pow(Fact[N], MOD - 2);

        for (valueType i = N - 1; i >= 0; --i)
            InvFact[i] = mul(InvFact[i + 1], i + 1);
    }

    valueType operator()(valueType n, valueType m) {
        if (n < 0 || m < 0 || n < m)
            return 0;

        if (m > N)
            throw std::out_of_range("BinomialCoefficient::operator() : m > N");

        if (n <= N)
            return mul(Fact[n], mul(InvFact[m], InvFact[n - m]));

        valueType result = 1;

        for (valueType i = 0; i < m; ++i)
            Mul(result, n - i);

        Mul(result, InvFact[m]);

        return result;
    }
};

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    std::cout.tie(nullptr);

    valueType T;

    std::cin >> T;

    for (valueType testcase = 1; testcase <= T; ++testcase) {
        valueType N, V;

        std::cin >> N >> V;

        BinomialCoefficient C(2 * N);

        ValueVector LeftSon(N + 1), RightSon(N + 1), Weight(N + 1);

        for (valueType i = 1; i <= N; ++i)
            std::cin >> LeftSon[i] >> RightSon[i] >> Weight[i];

        PairVector Limits;

        valueType dfsCount = 0;

        auto dfs = [&](auto self, valueType x) -> void {
            if (LeftSon[x] != -1)
                self(self, LeftSon[x]);

            ++dfsCount;

            if (Weight[x] != -1)
                Limits.emplace_back(dfsCount, Weight[x]);

            if (RightSon[x] != -1)
                self(self, RightSon[x]);
        };

        dfs(dfs, 1);

        Limits.emplace_back(0, 1);
        Limits.emplace_back(N + 1, V);

        std::sort(Limits.begin(), Limits.end());

        valueType ans = 1;

        for (valueType i = 1; i < Limits.size(); ++i)
            Mul(ans, C(Limits[i].first - Limits[i - 1].first - 1 + Limits[i].second - Limits[i - 1].second, Limits[i].first - Limits[i - 1].first - 1));

        std::cout << ans << std::endl;
    }

    return 0;
}